from __future__ import print_function

import copy
import math
import os
import sys
from collections import defaultdict
from operator import gt
from select import select
from time import sleep

import matplotlib.pyplot as plt
import numpy as np

from munkres import Munkres

# import matplotlib; matplotlib.use('Agg')

"""
Data format：
frame_id track_id obj_class  x y z

obj_class:
0 target_car
1 other_car
2 Pestarin
"""

def center_dis(gt,pre):
    delta_x = gt.x - pre.x
    delta_y = gt.y-pre.y
    delta_z = gt.z - pre.z
    return math.sqrt(delta_x*delta_x+ delta_y*delta_y+ delta_z*delta_z)

class tData:
    def __init__(self,frame=-1,obj_class="unset",\
                 x=-1000,y=-1000,z=-1000,track_id=-1):        
        # init object data
        self.frame      = frame
        self.track_id   = track_id
        self.obj_class   = obj_class
        self.x          = x
        self.y          = y
        self.z          = z

    def __str__(self):
        """
            Print read data.
        """
        attrs = vars(self)
        return '\n'.join("%s: %s" % item for item in attrs.items())

class trackingEvaluation(object):
    """ tracking statistics (CLEAR MOT, id-switches, fragments, ML/PT/MT, precision/recall)
             MOTA   - Multi-object tracking accuracy in [0,100]
             MOTP   - Multi-object tracking precision in [0,100] (3D) / [td,100] (2D)
    """

    def __init__(self, trk_result_path, gt_path, class_type):
        self.threhold = 1
        self.cls = class_type
        gt_file = open(gt_path, "r")
        gt_trks_list = gt_file.readlines()
        end_line_data= np.array(gt_trks_list[len(gt_trks_list)-1].strip().split(' '), dtype = np.float64 )
        self.num_frame = int(end_line_data[0])+1

        self.gt_trajectories            = [[] for x in range(self.num_frame)]
        self.trks_trajectories           = [[] for x in range(self.num_frame)]
        self.loadData(gt_path,True)
        self.loadData(trk_result_path,False)

        self.num_gt = 0
        self.true_positive = 0
        self.num_gt_trk_pair = 0
        self.true_positive_dis = []
        self.false_positive = 0
        self.num_miss = 0
        self.last_gt_trk_pair = []
        self.mismatched = 0

    def loadData(self,path,is_gt=True):
        file = open(path, "r")
        trks_list = file.readlines()
        for i in range(len(trks_list)):
            line_data= np.array(trks_list[i].strip().split(' '), dtype = np.float64 )
            if self.cls == "car" and line_data[2] == 2:
                continue
            if self.cls == "pedestrian" and line_data[2] <=1:
                continue
            t_data  = tData()
            t_data.frame = int(line_data[0])
            t_data.track_id = int(line_data[1])
            t_data.obj_class = int(line_data[2])
            t_data.x = line_data[3]
            t_data.y = line_data[4]
            t_data.z = line_data[5]
            if(is_gt):
                self.gt_trajectories[t_data.frame].append(t_data)
            else:
                self.trks_trajectories[t_data.frame].append(t_data)
    
    def analyse(self):
        for frame in  range(self.num_frame):
            self.num_gt += len(self.gt_trajectories[frame])
            # use hungarian method to associate, using 0..1 as cost
            # build cost matrix
            # row is gt, column is det
            cost_matrix = []
            this_ids_pair = [[],[]]
            for gg in self.gt_trajectories[frame]:
                # save current ids
                this_ids_pair[0].append(gg.track_id)
                this_ids_pair[1].append(-1)
                cost_row = []
                for tt in self.trks_trajectories[frame]:
                    c = center_dis(gg, tt)
                    # gating for center distance
                    if c <= self.threhold:
                        cost_row.append(c)
                    else:
                        cost_row.append(max_cost) # = 1e9
                cost_matrix.append(cost_row)

            if len(self.gt_trajectories[frame]) == 0:
                cost_matrix=[[]]
            # associate
            association_matrix = hm.compute(cost_matrix)

            # mapping for tracker ids and ground truth ids
            for row,col in association_matrix:
                # apply gating on box overlap
                c = cost_matrix[row][col]
                if c < max_cost:
                    this_ids_pair[1][row] = self.trks_trajectories[frame][col].track_id
                    # true positives are only valid associations
                    self.true_positive += 1
                    self.true_positive_dis.append(c)

                else:  #匹配错误 导致
                    self.false_positive+=1
                    self.num_miss +=1
            
            self.num_miss+= len(self.gt_trajectories[frame]) - len(association_matrix)  #检测少了的情况
            temp =len(self.trks_trajectories[frame]) -len(self.gt_trajectories[frame])  #检测多了的情况
            if(temp>0):
                self.false_positive+=temp
            
            #错误匹配
            if(len(self.last_gt_trk_pair)):  #第一帧没有跟踪值
                for i,tt in enumerate(this_ids_pair[0]):
                    if tt in self.last_gt_trk_pair[0]:
                        idx = self.last_gt_trk_pair[0].index(tt)
                        tid = this_ids_pair[1][i]            # id in current tracker corresponding to the gt tt
                        lid = self.last_gt_trk_pair[1][idx]          # id in last frame tracker corresponding to the gt tt
                        if tid != lid and lid != -1 and tid != -1:
                            self.mismatched+=1
                #不能直接将匹配对覆盖，因为当前跟踪丢失不代表之后还丢失，根据真值情况覆盖
                temp_pair = [[],[]]
                for i,tt in enumerate(this_ids_pair[0]):
                    if tt in self.last_gt_trk_pair[0]:
                        idx = self.last_gt_trk_pair[0].index(tt)
                        tid = this_ids_pair[1][i]            # id in current tracker corresponding to the gt tt
                        lid = self.last_gt_trk_pair[1][idx]          # id in last frame tracker corresponding to the gt tt
                        if tid == lid:
                            temp_pair[0].append(tt)
                            temp_pair[1].append(tid) 
                        if tid != lid and lid == -1 and tid != -1:
                            temp_pair[0].append(tt)
                            temp_pair[1].append(tid) 
                        if tid != lid and lid != -1 and tid == -1:
                            temp_pair[0].append(tt)
                            temp_pair[1].append(lid) 
                        if tid != lid and lid != -1 and tid != -1:
                            temp_pair[0].append(tt)
                            temp_pair[1].append(tid) 
                    else:
                        temp_pair[0].append(tt)
                        temp_pair[1].append(this_ids_pair[1][i])
                self.last_gt_trk_pair = temp_pair
            else:
                self.last_gt_trk_pair = this_ids_pair
            
        #MOTA
        mota = sum(self.true_positive_dis) / self.true_positive
        motp = 1 - float(self.num_miss)/self.num_gt \
            - float(self.false_positive)/self.num_gt \
                - float(self.mismatched)/self.num_gt
        
        print("MOTA: ",mota)
        print("MOTP: ",motp)
            

if __name__ == "__main__":
    trk_result_path = "/media/st/data_fast/gazebo/0813/envalute/testing/trajectory/trk_rejuct_tra_result.txt"
    # MOTA:  0.14406116811158912
    # MOTP:  0.7666912306558586
    # trk_result_path = "/media/st/data_fast/gazebo/0813/envalute/testing/trajectory/trk_without_dis_limit_result.txt"
    # MOTA:  0.14374149620334925
    # MOTP:  0.7799557848194547
    # trk_result_path = "/media/st/data_fast/gazebo/0813/envalute/testing/trajectory/trk_pc_trk_img_result.txt"
    gt_path = "/media/st/data_fast/gazebo/0813/envalute/testing/trajectory/tracking_label.txt"
    trk_en = trackingEvaluation(trk_result_path,gt_path,"all")
    trk_en.calMetric()
